{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "\n",
    "import loguru\n",
    "from torch.nn import functional as F\n",
    "\n",
    "from dataset import *\n",
    "from module import *\n",
    "%matplotlib inline\n",
    "torch.cuda.is_available()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define some constants\n",
    "TOKENIZE_MAX_LENGTH = 80\n",
    "WORD_VEC_DIM = 256"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# init dataset and other stuff\n",
    "# torch.manual_seed(42)\n",
    "root_dir = \"../deepfashion-multimodal/\"\n",
    "image_dir = path.join(root_dir, \"images\")\n",
    "train_dataset = DeepFashionDataset(path.join(root_dir, \"train_captions.json\"), True)\n",
    "test_dataset = DeepFashionDataset(path.join(root_dir, \"test_captions.json\"), False)\n",
    "loguru.logger.remove()\n",
    "loguru.logger.add(\"logs/logs_{time}.log\")\n",
    "log = loguru.logger"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# calculate some indicators about the dataset labels\n",
    "from numpy import mean\n",
    "\n",
    "\n",
    "def load_word_map(_dir=\"./word_map.json\", save=True, override=False):\n",
    "    if not override and os.path.exists(_dir):\n",
    "        with open(_dir, \"r\") as f:\n",
    "            _word2id = json.load(f)\n",
    "        _id2word = {v: k for k, v in _word2id.items()}\n",
    "\n",
    "    else:\n",
    "        labels = list(train_dataset.labels) + list(test_dataset.labels)\n",
    "        word = set()\n",
    "        sen_len = list()\n",
    "        for l in labels:\n",
    "            l = re.findall(r\"[\\w-]+|[^\\s\\w-]\", l)\n",
    "            sen_len.append(len(l))\n",
    "            word.update(l)\n",
    "        _word2id = {w: i for i, w in enumerate(word)}\n",
    "        _id2word = {i: w for i, w in enumerate(word)}\n",
    "\n",
    "        extra_token = [\"<pad>\", \"<start>\", \"<end>\", \"<unk>\"]\n",
    "        for w in extra_token:\n",
    "            word.add(w)\n",
    "            _word2id[w] = len(_word2id)\n",
    "            _id2word[len(_id2word)] = w\n",
    "\n",
    "        if save:\n",
    "            with open(_dir, \"w\") as f:\n",
    "                json.dump(_word2id, f)\n",
    "\n",
    "        _word_num = len(word)\n",
    "        print(\"word num: \", _word_num)\n",
    "        print(\"max sentence length: \", max(sen_len))\n",
    "        print(\"min sentence length: \", min(sen_len))\n",
    "        print(\"average sentence length: \", round(mean(sen_len), 2))\n",
    "        print(\"percentile 90 sentence len:\", np.percentile(sen_len, 90))\n",
    "    return _word2id, _id2word\n",
    "\n",
    "\n",
    "word2id, id2word = load_word_map(\"./word_map.json\", True, False)\n",
    "word_num = len(word2id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# init dataloader\n",
    "def tokenize(caption):\n",
    "    caption = re.findall(r\"[\\w-]+|[^\\s\\w-]\", caption)\n",
    "    caption = [word2id.get(w, word2id[\"<unk>\"]) for w in caption]\n",
    "    return caption\n",
    "\n",
    "\n",
    "def encode(caption, total_len=72):\n",
    "    caption = re.findall(r\"[\\w-]+|[^\\s\\w-]\", caption)\n",
    "    caption = [word2id.get(w, word2id[\"<unk>\"]) for w in caption]\n",
    "    caption = caption[: total_len - 2]\n",
    "    caption = [word2id[\"<start>\"]] + caption + [word2id[\"<end>\"]]\n",
    "    caption = caption + [word2id[\"<pad>\"]] * (total_len - len(caption))\n",
    "    return caption\n",
    "\n",
    "\n",
    "def decode(caption):\n",
    "    caption = [id2word.get(i, \"<unk>\") for i in caption]\n",
    "    return \" \".join(caption)\n",
    "\n",
    "\n",
    "def collect_fn(batch_data):\n",
    "    _images = []\n",
    "    _labels = []\n",
    "    # max_len = max(len(tokenize(label)) for _, label in batch_data)\n",
    "    # max_len = min(max_len, TOKENIZE_MAX_LENGTH)\n",
    "    for img, label_enc in batch_data:\n",
    "        # label_enc = encode(label_enc, total_len=max_len)\n",
    "        label_enc = encode(label_enc, total_len=TOKENIZE_MAX_LENGTH)\n",
    "        _images.append(img)\n",
    "        _labels.append(torch.tensor(label_enc))\n",
    "    _images = torch.stack(_images)\n",
    "    _labels = torch.stack(_labels)\n",
    "    return _images, _labels\n",
    "\n",
    "\n",
    "def get_dataloader(dataset, batch_size=32, shuffle=True, collate_fn=collect_fn, num_workers=4):\n",
    "    return DataLoader(\n",
    "        dataset, batch_size=batch_size, shuffle=shuffle, collate_fn=collate_fn, num_workers=num_workers\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# init resnet, transformer etc.\n",
    "def init_train(lr=1e-4, _resnet_lr=3e-5, _device=torch.device(\"cpu\"), **kwargs):\n",
    "    _resnet_transformer = MyTransformer(**kwargs).to(_device)\n",
    "    resnet_param = _resnet_transformer.resnet_train_params\n",
    "    other_param = _resnet_transformer.other_params\n",
    "\n",
    "    # optim = torch.optim.Adam(\n",
    "    #     [\n",
    "    #         {\"params\": other_param, \"lr\": lr},\n",
    "    #         {\"params\": resnet_param, \"lr\": _resnet_lr},\n",
    "    #     ]\n",
    "    # )\n",
    "    optim = torch.optim.AdamW(\n",
    "        [\n",
    "            {\"params\": other_param, \"lr\": lr, \"weight_decay\": 1e-4, \"amsgrad\": True},\n",
    "            {\"params\": resnet_param, \"lr\": _resnet_lr, \"weight_decay\": 5e-5, \"amsgrad\": False},\n",
    "        ]\n",
    "    )\n",
    "    # optim = torch.optim.SGD(\n",
    "    #     [\n",
    "    #         {\"params\": other_param, \"lr\": lr, \"weight_decay\": 1e-4, \"momentum\": 0.9},\n",
    "    #         {\"params\": resnet_param, \"lr\": _resnet_lr, \"weight_decay\": 1e-4, \"momentum\": 0.9},\n",
    "    #     ]\n",
    "    # )\n",
    "\n",
    "    log.info(\"init MyTransformer :\" + _resnet_transformer.train_indicator)\n",
    "    log.info(f\"init Adam lr: {lr}\")\n",
    "    log.info(\"device: {}\", _device)\n",
    "    print(\"use device: \", _device)\n",
    "    return _resnet_transformer, optim, _device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from meteor_util import safe_calculate_meteor_score\n",
    "def train(\n",
    "    model: MyTransformer, optim, _lr, _device, epoch=30, batch_size=32, num_workers=0, fine_tuning_range=1.0\n",
    "):\n",
    "    try:\n",
    "        model_path = f\"./model/b{batch_size}/e{epoch}/lr{_lr}\"\n",
    "        os.makedirs(model_path, exist_ok=True)\n",
    "        with open(f\"{model_path}/params.txt\", \"w\") as f:\n",
    "            f.write(model.train_indicator)\n",
    "        log.info(f\"start training, total epoch: {epoch}, batch_size: {batch_size}\")\n",
    "\n",
    "        finetuning_flag = False\n",
    "\n",
    "        for e in range(1, epoch + 1):\n",
    "            if not finetuning_flag and e > fine_tuning_range * epoch:\n",
    "                finetuning_flag = True\n",
    "                log.info(f\"start resnet fine tuning, epoch:{e}\")\n",
    "                model.partly_unfreeze_resnet()\n",
    "\n",
    "            _train = get_dataloader(train_dataset, batch_size=batch_size, num_workers=num_workers)\n",
    "            total_step = len(_train)\n",
    "            log.info(f\"epoch: {e}\")\n",
    "\n",
    "            for idx, (img, label) in (\n",
    "                bar := tqdm(\n",
    "                    enumerate(_train),\n",
    "                    total=total_step,\n",
    "                )\n",
    "            ):\n",
    "                img, label = img.to(_device), label.to(_device)\n",
    "                bar.set_description(f\"training epoch: {e}, step: {idx} \")\n",
    "                optim.zero_grad()\n",
    "\n",
    "                input_seq = label[:, :-1]\n",
    "                tgt_seq = label[:, 1:]\n",
    "                padding_mask = input_seq == word2id[\"<pad>\"]\n",
    "                tgt_mask = model.get_tgt_mask(TOKENIZE_MAX_LENGTH - 1, device=_device)\n",
    "\n",
    "                output = model.forward(\n",
    "                    img,\n",
    "                    input_seq,\n",
    "                    tgt_key_padding_mask=padding_mask,\n",
    "                    resnet_finetuning=finetuning_flag,\n",
    "                    tgt_mask=tgt_mask,\n",
    "                )\n",
    "                output, tgt_seq = output.flatten(0, 1), tgt_seq.flatten()\n",
    "                loss = F.cross_entropy(output, tgt_seq)\n",
    "                loss.backward()\n",
    "                nn.utils.clip_grad_norm_(model.parameters(), 5)\n",
    "                optim.step()\n",
    "\n",
    "                bar.set_postfix(loss=loss.item())\n",
    "                if idx % 10 == 0:\n",
    "                    log.info(f\"epoch:{e},step:{idx} loss: {loss.item()}\")\n",
    "\n",
    "            if e % 5 == 0:\n",
    "                test_loss = calculate_loss(model, test_dataset, _device=_device, batch_size=batch_size)\n",
    "                log.info(f\"epoch:{e}, test loss: {test_loss}\")\n",
    "                torch.save(\n",
    "                    {\"model\": model.state_dict(), \"optim\": optim.state_dict()},\n",
    "                    f\"{model_path}/model_epoch{e}_loss{test_loss.item()}.pth\",\n",
    "                )\n",
    "\n",
    "        log.info(\"training finished\")\n",
    "        if epoch % 5 != 0:\n",
    "            test_loss = calculate_loss(model, test_dataset, _device=_device, batch_size=batch_size)\n",
    "            torch.save(\n",
    "                {\"model\": model.state_dict(), \"optim\": optim.state_dict()},\n",
    "                f\"{model_path}/model_epoch{epoch}_loss{test_loss.item()}.pth\",\n",
    "            )\n",
    "    except Exception as e:\n",
    "        log.error(e)\n",
    "        raise e\n",
    "\n",
    "def calculate_loss(model, dataset, _device, batch_size=32, num_workers=0):\n",
    "    with torch.no_grad():\n",
    "        _test = get_dataloader(dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers)\n",
    "        losses = []\n",
    "        for idx, (img, label) in enumerate(_test):\n",
    "            img, label = img.to(_device), label.to(_device)\n",
    "            input_seq = label[:, :-1]\n",
    "            tgt_seq = label[:, 1:]\n",
    "            padding_mask = input_seq == word2id[\"<pad>\"]\n",
    "\n",
    "            output = model(img, input_seq, tgt_key_padding_mask=padding_mask)\n",
    "\n",
    "            output, tgt_seq = output.flatten(0, 1), tgt_seq.flatten()\n",
    "            loss = F.cross_entropy(output, tgt_seq)\n",
    "            losses.append(loss.item())\n",
    "        return round(mean(losses), 5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# init resnet, transformer etc.\n",
    "\n",
    "# BATCH_SIZE = 128\n",
    "BATCH_SIZE = 28\n",
    "learning_rate =6.5e-5\n",
    "resnet_lr = 5e-5\n",
    "\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "resnet_transformer, optimizer, device = init_train(\n",
    "    lr=learning_rate,\n",
    "    _resnet_lr=resnet_lr,\n",
    "    word_num=word_num,\n",
    "    word_vec_dim=WORD_VEC_DIM,\n",
    "    _device=device,\n",
    "    num_decoder_layers=12,\n",
    "    num_encoder_layers=12,\n",
    "    activation=\"gelu\",\n",
    "    dropout=0.3,\n",
    "    temperature=1000,\n",
    "    nhead=8,\n",
    "    dim_feedforward=2048,\n",
    "    custom_encoder=None,\n",
    "    custom_decoder=None,\n",
    "    layer_norm_eps=1e-5,\n",
    ")\n",
    "# model = \"./model/b128/b40/lr0.001/imagenet_model_epoch10_loss0.40721.pth\"\n",
    "# resnet_transformer.load_state_dict(\n",
    "#     torch.load(model)[\"model\"]\n",
    "# )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train(\n",
    "    resnet_transformer,\n",
    "    optimizer,\n",
    "    epoch=30,\n",
    "    _lr=learning_rate,\n",
    "    _device=device,\n",
    "    batch_size=BATCH_SIZE,\n",
    "    num_workers=0,\n",
    "    fine_tuning_range=0.6,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ML",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
